import datetime as dt
import time
import requests
from zoneinfo import ZoneInfo
import polars as pl
from tqdm import tqdm
from functools import partial
from pathlib import Path
from concurrent.futures import ThreadPoolExecutor

from cryptoAlpha.dataFeed.IDataFeed import IDataFeed

CHINA_TZ: ZoneInfo = ZoneInfo("Asia/Shanghai")

def parse_timestamp(timestamp: int) -> dt.datetime:
    datetime = dt.datetime.fromtimestamp(int(timestamp) / 1000)
    return datetime.replace(tzinfo=CHINA_TZ)

class OkeDataFeed(IDataFeed):
    def __init__(self) -> None:
        super().__init__("OkeDataFeed")
        self.exchange = "OKE"
        self.base_url = "https://www.okx.com"
        
    def get_contract_info(self, inst_type):
        """查询合约"""
        path = self.base_url + "/api/v5/public/instruments"
        resp = requests.get(
            path,
            params={"instType": inst_type}
        )
        return resp.json()

    def get_kline(self, symbol: str, freq: str, start_time: dt.datetime=None, end_time: dt.datetime=None):
        """获取K线数据"""
        path = self.base_url + "/api/v5/market/history-candles"
        
        start_time_stamp = int(start_time.timestamp() * 1000) if start_time else None
        end_time_stamp = int(end_time.timestamp() * 1000) if end_time else None
        cursor = None
        df = pl.DataFrame()

        params = {
            "instId": symbol,
            "bar": freq,
        }
        pbar = tqdm()
        # K线数据每个粒度最多可获取最近1440条，这里循环获取
        while True:
            pbar.set_description(f"{symbol}: {parse_timestamp(start_time_stamp) if start_time_stamp else ''} - {parse_timestamp(end_time_stamp) if end_time_stamp else ''}")
            pbar.update(1)
            if start_time_stamp and end_time_stamp <= start_time_stamp:
                    break
            if start_time_stamp:
                params["before"] = start_time_stamp
            if end_time_stamp:
                params["after"] = end_time_stamp

            data = None
            for i in range(10):
                try:
                    # 从服务器获取响应
                    resp = requests.get(
                        path,
                        params=params
                    )
                                # 如果请求失败则终止循环
                    if resp.status_code // 100 != 2:
                        msg = f"获取历史数据失败，状态码：{resp.status_code}，信息：{resp.text}"
                        time.sleep(10)
                    else:
                        data: dict = resp.json()
                        break
                except Exception as e:
                    msg = f"获取历史数据失败，错误：{e}"
                    print(msg)
                    time.sleep(10)
                
            if not data:
                msg = f"获取历史数据失败"
                break
            if not data["data"]:
                msg = f"获取历史数据为空，{params} {data['msg']}"
                print(msg)
                break

            for bar_list in data["data"]:
                ts, o, h, l, c, vol, volcc, volCcyQuote, confirm = bar_list
                cursor = int(ts) if not cursor else min(cursor, int(ts))
                if start_time_stamp and cursor < start_time_stamp:
                    break
                if confirm == 0:
                    continue
                bar = pl.DataFrame({
                    "trading_date": parse_timestamp(ts).date(),
                    "datetime" : parse_timestamp(ts),
                    "open" : float(o),
                    "high" : float(h),
                    "low" : float(l),
                    "close" : float(c),
                    "volume" : float(vol)
                })
                df = df.vstack(bar)
                
                end_time_stamp = cursor
        return df

if __name__ == "__main__":

    symbol = "BTC-USDT"
    freq = "1m"
    start_time = dt.datetime(2020,1,1)
    end_time = dt.datetime(2025,1,1)
    # path = "https://www.okx.com" + "/api/v5/market/history-candles"
    feed = OkeDataFeed()
    
    contract = feed.get_contract_info("SWAP")
    dataPath = Path("data").joinpath("M1")
    dataPath.mkdir(parents=True, exist_ok=True)
    # 创建线程池
    with ThreadPoolExecutor(max_workers=5) as executor:
        ed = end_time
        while ed >= start_time:
            sd = ed-dt.timedelta(days=5)
            for c in tqdm(contract["data"]):
                symbol = c["instId"]
                task = executor.submit(feed.get_kline, symbol, freq, sd, ed)
                task.add_done_callback(lambda future: future.result().write_parquet(dataPath.joinpath(f"OKX_{symbol}_{freq}.parquet").as_posix(), partition_by="trading_date"))
            ed = sd